//
// Created by ZhaoXiaoFei on 2022/6/28.
//

#ifndef VINS_ESTIMATOR_PREINTEGRATION_HPP
#define VINS_ESTIMATOR_PREINTEGRATION_HPP

#include <ceres/ceres.h>
#include <eigen3/Eigen/Dense>
#include <sophus/so3.hpp>
#include <iostream>
#include "../utility/utility.h"
#include <glog/logging.h>
using namespace Eigen;

class IMUpreintegration{
public:
    IMUpreintegration() = default;
    IMUpreintegration(const Eigen::Vector3d& acc_0_, const Eigen::Vector3d& gyr_0_,
                      const Eigen::Vector3d& linearized_ba_, const Eigen::Vector3d& linearized_bg_)
            : acc_0(acc_0_), gyr_0(gyr_0_),
              linearized_acc(acc_0_), linearized_gyr(gyr_0_),
              linearized_ba(linearized_ba_), linearized_bg(linearized_bg_){
        sum_dt = 0.0;
        deltaP = Eigen::Vector3d::Zero();
        deltaQ = Eigen::Quaterniond::Identity();
        deltaV = Eigen::Vector3d::Zero();

        jacobian = Eigen::Matrix<double, 15, 15>::Identity();
        covariance = Eigen::Matrix<double, 15, 15>::Zero();

        noise = Eigen::Matrix<double, 18, 18>::Zero();
        noise.block<3,3>(0,0) = ACC_N * ACC_N * Eigen::Matrix3d::Identity();
        noise.block<3,3>(3,3) = GYR_N * GYR_N * Eigen::Matrix3d::Identity();
        noise.block<3,3>(6,6) = ACC_N * ACC_N * Eigen::Matrix3d::Identity();
        noise.block<3,3>(9,9) = GYR_N * GYR_N * Eigen::Matrix3d::Identity();
        noise.block<3,3>(12,12) = ACC_W * ACC_W * Eigen::Matrix3d::Identity();
        noise.block<3,3>(15,15) = GYR_W * GYR_W * Eigen::Matrix3d::Identity();

    }
    void push_back(double dt, const Eigen::Vector3d& acc, const Eigen::Vector3d& gyr){
        dt_buf.push_back(dt);
        acc_buf.push_back(acc);
        gyr_buf.push_back(gyr);
        propagate(dt, acc, gyr);
    }
    void repropagate(const Eigen::Vector3d& ba, const Eigen::Vector3d& bg){
        sum_dt = 0.0;
        acc_0 = linearized_acc;
        gyr_0 = linearized_gyr;
        linearized_ba = ba;
        linearized_bg = bg;
        deltaP.setZero();
        deltaQ.setIdentity();
        deltaV.setZero();
        jacobian.setIdentity();
        covariance.setZero();

        for(int i = 0; i < dt_buf.size(); i++){
            propagate(dt_buf[i], acc_buf[i], gyr_buf[i]);
        }
    }
    void midPointIntegration(double _dt,
                             const Eigen::Vector3d &_acc_0, const Eigen::Vector3d &_gyr_0,
                             const Eigen::Vector3d &_acc_1, const Eigen::Vector3d &_gyr_1,
                             const Eigen::Vector3d &delta_p, const Eigen::Quaterniond &delta_q, const Eigen::Vector3d &delta_v,
                             const Eigen::Vector3d &linearized_ba, const Eigen::Vector3d &linearized_bg,
                             Eigen::Vector3d &result_delta_p, Eigen::Quaterniond &result_delta_q, Eigen::Vector3d &result_delta_v,
                             Eigen::Vector3d &result_linearized_ba, Eigen::Vector3d &result_linearized_bg, bool update_jacobian)
    {
        //ROS_INFO("midpoint integration");
        Vector3d un_acc_0 = delta_q * (_acc_0 - linearized_ba);
        Vector3d un_gyr = 0.5 * (_gyr_0 + _gyr_1) - linearized_bg;
        result_delta_q = delta_q * Quaterniond(1, un_gyr(0) * _dt / 2, un_gyr(1) * _dt / 2, un_gyr(2) * _dt / 2);
        Vector3d un_acc_1 = result_delta_q * (_acc_1 - linearized_ba);
        Vector3d un_acc = 0.5 * (un_acc_0 + un_acc_1);
        result_delta_p = delta_p + delta_v * _dt + 0.5 * un_acc * _dt * _dt;
        result_delta_v = delta_v + un_acc * _dt;
        result_linearized_ba = linearized_ba;
        result_linearized_bg = linearized_bg;

        if(update_jacobian)
        {
            Vector3d w_x = 0.5 * (_gyr_0 + _gyr_1) - linearized_bg;
            Vector3d a_0_x = _acc_0 - linearized_ba;
            Vector3d a_1_x = _acc_1 - linearized_ba;
            Matrix3d R_w_x, R_a_0_x, R_a_1_x;

            R_w_x<<0, -w_x(2), w_x(1),
                    w_x(2), 0, -w_x(0),
                    -w_x(1), w_x(0), 0;
            R_a_0_x<<0, -a_0_x(2), a_0_x(1),
                    a_0_x(2), 0, -a_0_x(0),
                    -a_0_x(1), a_0_x(0), 0;
            R_a_1_x<<0, -a_1_x(2), a_1_x(1),
                    a_1_x(2), 0, -a_1_x(0),
                    -a_1_x(1), a_1_x(0), 0;

            MatrixXd F = MatrixXd::Zero(15, 15);
            F.block<3, 3>(0, 0) = Matrix3d::Identity();
            F.block<3, 3>(0, 3) = -0.25 * delta_q.toRotationMatrix() * R_a_0_x * _dt * _dt +
                                  -0.25 * result_delta_q.toRotationMatrix() * R_a_1_x * (Matrix3d::Identity() - R_w_x * _dt) * _dt * _dt;
            F.block<3, 3>(0, 6) = MatrixXd::Identity(3,3) * _dt;
            F.block<3, 3>(0, 9) = -0.25 * (delta_q.toRotationMatrix() + result_delta_q.toRotationMatrix()) * _dt * _dt;
            F.block<3, 3>(0, 12) = -0.25 * result_delta_q.toRotationMatrix() * R_a_1_x * _dt * _dt * -_dt;
            F.block<3, 3>(3, 3) = Matrix3d::Identity() - R_w_x * _dt;
            F.block<3, 3>(3, 12) = -1.0 * MatrixXd::Identity(3,3) * _dt;
            F.block<3, 3>(6, 3) = -0.5 * delta_q.toRotationMatrix() * R_a_0_x * _dt +
                                  -0.5 * result_delta_q.toRotationMatrix() * R_a_1_x * (Matrix3d::Identity() - R_w_x * _dt) * _dt;
            F.block<3, 3>(6, 6) = Matrix3d::Identity();
            F.block<3, 3>(6, 9) = -0.5 * (delta_q.toRotationMatrix() + result_delta_q.toRotationMatrix()) * _dt;
            F.block<3, 3>(6, 12) = -0.5 * result_delta_q.toRotationMatrix() * R_a_1_x * _dt * -_dt;
            F.block<3, 3>(9, 9) = Matrix3d::Identity();
            F.block<3, 3>(12, 12) = Matrix3d::Identity();
            //cout<<"A"<<endl<<A<<endl;

            MatrixXd V = MatrixXd::Zero(15,18);
            V.block<3, 3>(0, 0) =  0.25 * delta_q.toRotationMatrix() * _dt * _dt;
            V.block<3, 3>(0, 3) =  0.25 * -result_delta_q.toRotationMatrix() * R_a_1_x  * _dt * _dt * 0.5 * _dt;
            V.block<3, 3>(0, 6) =  0.25 * result_delta_q.toRotationMatrix() * _dt * _dt;
            V.block<3, 3>(0, 9) =  V.block<3, 3>(0, 3);
            V.block<3, 3>(3, 3) =  0.5 * MatrixXd::Identity(3,3) * _dt;
            V.block<3, 3>(3, 9) =  0.5 * MatrixXd::Identity(3,3) * _dt;
            V.block<3, 3>(6, 0) =  0.5 * delta_q.toRotationMatrix() * _dt;
            V.block<3, 3>(6, 3) =  0.5 * -result_delta_q.toRotationMatrix() * R_a_1_x  * _dt * 0.5 * _dt;
            V.block<3, 3>(6, 6) =  0.5 * result_delta_q.toRotationMatrix() * _dt;
            V.block<3, 3>(6, 9) =  V.block<3, 3>(6, 3);
            V.block<3, 3>(9, 12) = MatrixXd::Identity(3,3) * _dt;
            V.block<3, 3>(12, 15) = MatrixXd::Identity(3,3) * _dt;

            //step_jacobian = F;
            //step_V = V;
            jacobian = F * jacobian;
            covariance = F * covariance * F.transpose() + V * noise * V.transpose();
        }

    }

    void propagate(double dt_, const Eigen::Vector3d& acc_, const Eigen::Vector3d& gyr_){
        dt = dt_;
        acc_1 = acc_;
        gyr_1 = gyr_;

        Vector3d result_delta_p;
        Quaterniond result_delta_q;
        Vector3d result_delta_v;
        Vector3d result_linearized_ba;
        Vector3d result_linearized_bg;

        midPointIntegration(dt, acc_0, gyr_0, acc_1, gyr_1, deltaP, deltaQ, deltaV,
                            linearized_ba, linearized_bg,
                            result_delta_p, result_delta_q, result_delta_v,
                            result_linearized_ba, result_linearized_bg, 1);

//        Eigen::Vector3d un_acc_0 = deltaQ * (acc_0 - linearized_ba);
//        Eigen::Vector3d un_gyr = 0.5 * ((gyr_0 - linearized_bg) + (gyr_1 - linearized_bg));
//        Eigen::Quaterniond result_deltaQ = deltaQ * Eigen::Quaterniond(1, 0.5 * un_gyr[0] * dt, 0.5 * un_gyr[1] * dt, 0.5 * un_gyr[2] * dt);
//
//        Eigen::Vector3d un_acc_1 = result_deltaQ * (acc_1 - linearized_ba);
//        Eigen::Vector3d un_acc = 0.5 * (un_acc_0 + un_acc_1);
//
//        deltaP = deltaP + deltaV * dt + 0.5 * un_acc * dt * dt;
//        deltaV = deltaV + un_acc * dt;

        //update jacobian
//        Eigen::Vector3d a_0 = acc_0 - linearized_ba;
//        Eigen::Vector3d a_1 = acc_1 - linearized_ba;
//
//        Eigen::Matrix3d a_0_x, a_1_x, w_x;
//        a_0_x = Sophus::SO3d::hat(a_0);
//        a_1_x = Sophus::SO3d::hat(a_1);
//        w_x = Sophus::SO3d::hat(un_gyr);
//
//        //Eigen::Vector3d left = deltaQ *
//        Eigen::Matrix<double, 15, 15> F = Eigen::Matrix<double, 15, 15>::Zero();
//        F.block<3,3>(0,0) = Eigen::Matrix3d::Identity();
//        F.block<3,3>(0,3) = -0.25 * dt * dt * (deltaQ.toRotationMatrix() * a_0_x + result_deltaQ.toRotationMatrix() * a_1_x * (Eigen::Matrix3d::Identity() - w_x * dt));
//        F.block<3,3>(0,6) = Eigen::Matrix3d::Identity() * dt;
//        F.block<3,3>(0,9) = -0.25 * dt * dt * (deltaQ.toRotationMatrix() + result_deltaQ.toRotationMatrix());
//        F.block<3,3>(0,12) = 0.25 * dt * dt * dt * (result_deltaQ.toRotationMatrix() * a_1_x);
//        F.block<3,3>(3,3) = Eigen::Matrix3d::Identity() - w_x * dt;
//        F.block<3,3>(3,12) = -1 * Eigen::Matrix3d::Identity() * dt;
//
//        F.block<3,3>(6,3) = -0.5 * dt * (deltaQ.toRotationMatrix() * a_0_x + result_deltaQ.toRotationMatrix() * a_1_x * (Eigen::Matrix3d::Identity() - w_x * dt));
//        F.block<3,3>(6,6) = Eigen::Matrix3d::Identity();
//        F.block<3,3>(6,9) = -0.5 * dt * (deltaQ.toRotationMatrix() + result_deltaQ.toRotationMatrix());
//        F.block<3,3>(6,12) = 0.5 * dt * dt * result_deltaQ.toRotationMatrix() * a_1_x;
//        F.block<3,3>(9,9) = Eigen::Matrix3d::Identity();
//        F.block<3,3>(12,12) = Eigen::Matrix3d::Identity();
//
//        Eigen::Matrix<double, 15, 18> V = Eigen::Matrix<double, 15, 18>::Zero();
//        V.block<3,3>(0,0) = 0.25 * deltaQ.toRotationMatrix() * dt * dt;
//        V.block<3,3>(0,3) = -0.25 * dt * dt * 0.5 * dt * result_deltaQ.toRotationMatrix() * a_1_x;
//        V.block<3,3>(0,6) = 0.25 * result_deltaQ.toRotationMatrix() * dt * dt;
//        V.block<3,3>(0,9) = -0.25 * dt * dt * 0.5 * dt * result_deltaQ.toRotationMatrix() * a_1_x;
//        V.block<3,3>(3,3) = 0.5 * Eigen::Matrix3d::Identity() * dt;
//        V.block<3,3>(3,9) = 0.5 * Eigen::Matrix3d::Identity() * dt;
//        V.block<3,3>(6,0) = 0.5 * deltaQ.toRotationMatrix() * dt;
//        V.block<3,3>(6,3) = -0.5 * dt * result_deltaQ.toRotationMatrix() * a_1_x * 0.5 * dt;
//        V.block<3,3>(6,6) = 0.5 * result_deltaQ.toRotationMatrix() * dt;
//        V.block<3,3>(6,9) = -0.5 * dt * result_deltaQ.toRotationMatrix() * a_1_x * 0.5 * dt;
//        V.block<3,3>(9,12) = Eigen::Matrix3d::Identity() * dt;
//        V.block<3,3>(12,15) = Eigen::Matrix3d::Identity() * dt;
//
//        jacobian = F * jacobian;
//        //std::cout << "jacobian111: " << jacobian.block<3,3>(3, 12) << std::endl;
//        covariance = F * covariance * F.transpose() + V * noise * V.transpose();

//        deltaQ = result_deltaQ;
//        deltaQ.normalize();
//        sum_dt += dt;
//        acc_0 = acc_1;
//        gyr_0 = gyr_1;

        deltaP = result_delta_p;
        deltaQ = result_delta_q;
        deltaV = result_delta_v;
        linearized_ba = result_linearized_ba;
        linearized_bg = result_linearized_bg;
        deltaQ.normalize();
        sum_dt += dt;
        acc_0 = acc_1;
        gyr_0 = gyr_1;
    }

    Eigen::Matrix<double, 15, 1> evaluate(const Eigen::Vector3d& Pi, const Eigen::Quaterniond& Qi, const Eigen::Vector3d& Vi, const Eigen::Vector3d& Bai, const Eigen::Vector3d& Bgi,
                                          const Eigen::Vector3d& Pj, const Eigen::Quaterniond& Qj, const Eigen::Vector3d& Vj, const Eigen::Vector3d& Baj, const Eigen::Vector3d& Bgj){

        Eigen::Matrix3d dp_dba = jacobian.block<3,3>(0,9);
        Eigen::Matrix3d dp_dbg = jacobian.block<3,3>(0,12);

        Eigen::Matrix3d dq_dbg = jacobian.block<3,3>(3,12);

        Eigen::Matrix3d dv_dba = jacobian.block<3,3>(6,9);
        Eigen::Matrix3d dv_dbg = jacobian.block<3,3>(6,12);

        Eigen::Vector3d dba = Bai - linearized_ba;
        Eigen::Vector3d dbg = Bgi - linearized_bg;
//        LOG(INFO) << " Bgi: " << Bgi.transpose() << std::endl;
//        LOG(INFO) << " linearized_bg: " << linearized_bg.transpose() << std::endl;

        Eigen::Quaterniond correctedQ = deltaQ * Eigen::Quaterniond(1, (0.5 * dq_dbg * dbg)(0), (0.5 * dq_dbg * dbg)(1), (0.5 * dq_dbg * dbg)(2));
        Eigen::Vector3d correctedP = deltaP + dp_dba * dba + dp_dbg * dbg;
        Eigen::Vector3d correctedV = deltaV + dv_dba * dba + dv_dbg * dbg;

        Eigen::Matrix<double, 15, 1> residuals;

        residuals.block<3, 1>(0, 0) = Qi.inverse() * (Pj - Pi - Vi * sum_dt + 0.5 * G * sum_dt * sum_dt) - correctedP;
        residuals.block<3, 1>(3, 0) = 2 * (correctedQ.inverse() * Qi.inverse() * Qj).vec();
        residuals.block<3, 1>(6, 0) = Qi.inverse() * (G * sum_dt + Vj - Vi) - correctedV;
        residuals.block<3, 1>(9, 0) = Baj - Bai;
        residuals.block<3, 1>(12, 0) = Bgj - Bgi;
//        LOG(INFO) << " RESIDUAL1: " << residuals.block<3, 1>(9, 0).transpose() << std::endl;
//        LOG(INFO) << " RESIDUAL2: " << residuals.block<3, 1>(12, 0).transpose() << std::endl;
//        LOG(INFO) << " Bai: " << Bai.transpose() << " Baj: " << Baj.transpose() << std::endl;
//        LOG(INFO) << " Bgi: " << Bgi.transpose() << " Bgj: " << Bgj.transpose() << std::endl;

        return residuals;
    }



    double dt;
    Eigen::Vector3d acc_0, gyr_0;
    Eigen::Vector3d acc_1, gyr_1;

    const Eigen::Vector3d linearized_acc, linearized_gyr;
    Eigen::Vector3d linearized_ba, linearized_bg;

    Eigen::Matrix<double, 15, 15> jacobian, covariance;
    Eigen::Matrix<double, 18, 18> noise;

    double sum_dt;
    Eigen::Vector3d deltaP;
    Eigen::Quaterniond deltaQ;
    Eigen::Vector3d deltaV;


    //buf
    std::vector<double> dt_buf;
    std::vector<Eigen::Vector3d> acc_buf;
    std::vector<Eigen::Vector3d> gyr_buf;

};

#endif //VINS_ESTIMATOR_PREINTEGRATION_HPP
